package us.luosl.lslt.function.stream;

import us.luosl.lslt.collection.AbstractIterator;
import us.luosl.lslt.lang.Tuple;
import us.luosl.lslt.lang.Tuple2;

import java.util.*;
import java.util.function.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public interface GeneralStreamWrapper<T> extends StreamWrapper<T> {

    default <U> GeneralStreamWrapper<U> wrap(Stream<U> stream){
        return () -> stream;
    }

    default GeneralStreamWrapper<T> filter(Predicate<? super T> predicate) {
        return wrap(getStream().filter(predicate));
    }

    default <R> GeneralStreamWrapper<R> map(Function<? super T, ? extends R> mapper) {
        return wrap(getStream().map(mapper));
    }

    default <R> GeneralStreamWrapper<R>flatMap(Function<? super T, ? extends GeneralStreamWrapper<? extends R>> mapper) {
        return wrap(getStream().flatMap(data -> mapper.apply(data).getStream()));
    }

    default GeneralStreamWrapper<T> distinct() {
        return wrap(getStream().distinct());
    }

    default GeneralStreamWrapper<T> sorted() {
        return wrap(getStream().sorted());
    }

    default GeneralStreamWrapper<T> sorted(Comparator<? super T> comparator) {
        return wrap(getStream().sorted(comparator));
    }

    default GeneralStreamWrapper<T> peek(Consumer<? super T> action) {
        return wrap(getStream().peek(action));
    }

    default GeneralStreamWrapper<T> limit(long maxSize) {
        return wrap(getStream().limit(maxSize));
    }

    default GeneralStreamWrapper<T> skip(long n) {
        return wrap(getStream().skip(n));
    }

    // NEW API

    /**
     * 为 stream 加上当前索引，并且 StreamWrapper 会变更为 EntryStreamWrapper<T, Integer>
     * 例如:
     * GeneralStreamWrapper<String> stream = StreamWrapper.of("a", "b", "c");
     * stream.zipWithIndex().forEach((item, idx) -> System.out.println(idx + ":" + item));
     * 输出:
     * 0:a
     * 1:b
     * 2:c
     * @return 添加索引后的 EntryStreamWrapper
     */
    default EntryStreamWrapper<T, Integer> zipWithIndex() {
        Iterator<T> itr = iterator();
        AbstractIterator<Tuple2<T, Integer>> indexedItr;
        if(!itr.hasNext()){
            indexedItr = AbstractIterator.empty();
        }else{
            indexedItr = new AbstractIterator<Tuple2<T, Integer>>() {
                private int index = 0;

                @Override
                public boolean hasNext() {
                    return itr.hasNext();
                }

                @Override
                public Tuple2<T, Integer> getNext() {
                    return Tuple.of(itr.next(), index++);
                }
            };
        }

        return StreamWrapper.wrapEntry(indexedItr.stream());
    }

    default <T1, T2> EntryStreamWrapper<T1, T2> mapToEntry(Function<T, Tuple2<T1, T2>> mapper){
        return StreamWrapper.wrapEntry(getStream().map(mapper));
    }

    default GeneralStreamWrapper<T> sortBy(Supplier<Comparator<T>> supplier, boolean nullFirst){
        Comparator<T> comparator = supplier.get();
        if(nullFirst){
            comparator = Comparator.nullsFirst(comparator);
        }else{
            comparator = Comparator.nullsLast(comparator);
        }
        return sorted(comparator);
    }

    /**
     * 根据指定方式进行排序
     * 例如:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4, null);
     * List<Integer> s = stream.sortBy(i -> i, true).toList();
     * System.out.println(s);
     * 输出:
     * [null, 1, 2, 3, 4]
     * @param orderExtractor 排序字段抽取
     * @param nullFirst null 值是否放在头部
     * @param <R> 排序字段类型
     * @return 排序后的 streamWrapper
     */
    default <R extends Comparable<R>> GeneralStreamWrapper<T> sortBy(Function<T, R> orderExtractor, boolean nullFirst) {
       return sortBy(() -> Comparator.comparing(orderExtractor), nullFirst);
    }

    default <R extends Comparable<R>> GeneralStreamWrapper<T> sortBy(Function<T, R> orderExtractor) {
        return sortBy(() -> Comparator.comparing(orderExtractor), true);
    }

    default <R> GeneralStreamWrapper<T> sortBy(Function<T, R> orderExtractor, Comparator<R> byComparator, boolean nullFirst) {
        return sortBy(() -> Comparator.comparing(orderExtractor, byComparator), nullFirst);
    }

    /**
     * 根据指定方式进行倒排序
     * 例如:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4, null);
     * List<Integer> s = stream.sortBy(i -> i, true).toList();
     * System.out.println(s);
     * 输出:
     * [null, 4, 3, 2, 1]
     * @param orderExtractor 排序字段抽取
     * @param nullFirst null 值是否放在头部
     * @param <R> 排序字段类型
     * @return 排序后的 streamWrapper
     */
    default <R extends Comparable<R>> GeneralStreamWrapper<T> sortDescBy(Function<T, R> orderExtractor, boolean nullFirst){
        return sortBy(() -> Comparator.comparing(orderExtractor).reversed(), nullFirst);
    }

    default <R extends Comparable<R>> GeneralStreamWrapper<T> sortDescBy(Function<T, R> orderExtractor){
        return sortBy(() -> Comparator.comparing(orderExtractor).reversed(), false);
    }

    default <R> GeneralStreamWrapper<T> sortDescBy(Function<T, R> orderExtractor, Comparator<R> byComparator, boolean nullFirst) {
        return sortBy(() -> Comparator.comparing(orderExtractor, byComparator).reversed(), nullFirst);
    }

    /**
     * 反转该 StreamWrapper；但是反转代价较大
     * 例如：
     * Stream<String> s = Stream.of("4", "5", "1");
     * StreamWrapper.wrap(s).reverse().forEach(System.out::println)
     * 输出: 1,5,4
     * @return 反转后的 StreamWrapper
     */
    default GeneralStreamWrapper<T> reverse() {
        Iterator<T> itr = collect(Collectors.toCollection(LinkedList::new)).descendingIterator();
        return StreamWrapper.wrap(itr);
    }

    default <U> U foldRight(U zero, BiFunction<U, T, U> f) {
        return reverse().foldLeft(zero, f);
    }

    /**
     * 根据指定条件对 stream 进行去重
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 2, 4, null);
     * List<Integer> r = stream.distinctBy(i -> i).toList();
     * System.out.println(r);
     * 输出:
     * [2, 1, 4, null]
     * @param byExtractor 去重字段抽取
     * @param <K> 去重字段类型
     * @return 去重后的 streamWrapper
     */
    default <K> GeneralStreamWrapper<T> distinctBy(Function<T, K> byExtractor){
        final Set<K> distSet = new HashSet<>();
        return wrap(getStream().filter(item -> distSet.add(byExtractor.apply(item))));
    }

    default GeneralStreamWrapper<T> filterNot(Predicate<? super T> predicate){
        return filter(predicate.negate());
    }
}
